#include <swarm_robot_control.h>
#include <Hungary.h>

using namespace std;


/************************** Global Parameter ************************/ 
/* Convergence threshold */
double conv_x = 0.4; // Threshold of x, in m
double conv_y = 0.4;  // Threshold of y, in m
double conv_dis = 0.4;

/* Param for swarm */
std::vector<int> swarm_robot_id{1, 2, 3, 4, 5, 6, 7};       //Set ids of swarm robot based on Aruco marker
int num_robot = 5;                                          //the number of moving robot
int num_obs = 2;                                            //the number of obstacle
double swarm_vel = 0.1;                                     // the velocity of swarm, in m/s
double interval = 0.5;                                      //the distance within formation
int num_formation = 3;                                      //the number of formation
/*Formation type: 0:circle 1:line 2:cross 3:adaptcross 4:delta  
 *Note: the first formation in the list is init formation
 */
std::vector<SwarmRobot::FORMATION_TYPE> formation_name_list = {SwarmRobot::FORMATION_TYPE::CROSS, 
                                                               SwarmRobot::FORMATION_TYPE::LINE, 
                                                               SwarmRobot::FORMATION_TYPE::ADAPTCROSS};
// std::vector<SwarmRobot::FORMATION_TYPE> formation_name_list = {SwarmRobot::FORMATION_TYPE::CIRCLE, 
//                                                                SwarmRobot::FORMATION_TYPE::LINE, 
//                                                                SwarmRobot::FORMATION_TYPE::CROSS,
//                                                                SwarmRobot::FORMATION_TYPE::DELTA};

/* Velocity scale and threshold */
double MAX_W = 1;       // Maximum angle velocity (rad/s)
double MIN_W = 0.05;    // Minimum angle velocity(rad/s)
double MAX_V = 0.2;     // Maximum linear velocity(m/s)
double MIN_V = 0.0;     // Minimum linear velocity(m/s)
double k_w = 0.1;       // Scale of angle velocity
double k_v = 0.05;      // Scale of linear velocity

/* Parameters for Collision Avoidence */
double weight_repel = 1.0;
double weight_attract = 1.0;
double dis_repel = 0.4;
double dis_attract = 1.0; 


/************************** Function Statement ************************/ 
void getParameterList(ros::NodeHandle nh);


/************************** Main Function ************************/
int main(int argc, char** argv) {
    ros::init(argc, argv, "swarm_robot_control_formation");
    ros::NodeHandle nh("~");
    getParameterList(nh);

    /*Hungary*/
    Hungary hungary_swarm;
    hungary_swarm.num_swarm = num_robot;

    /*Formation to be choiced*/
    std::vector<Eigen::VectorXd> formation_x_list(num_formation);
    std::vector<Eigen::VectorXd> formation_y_list(num_formation);
    

    /* Set L Matrix */
    Eigen::MatrixXd lap(num_robot, num_robot);
    for (int i = 0; i < num_robot;i++){
        for (int j = 0; j < num_robot;j++){
            if(i==j){
                lap(i,j) = num_robot - 1;
            }
            else{
                lap(i, j) = -1;
            }
        }
    }

    /* Initialize swarm robot */
    SwarmRobot swarm_robot(&nh, swarm_robot_id);
    swarm_robot.MoveRobot_num = num_robot;
    swarm_robot.Obs_num = num_obs;
    swarm_robot.Formation_num = num_formation;
    swarm_robot.formation_name_list = formation_name_list;
    swarm_robot.lap = lap;

    /* Set Shape */
    Eigen::VectorXd final_target_x(num_robot);
    Eigen::VectorXd final_target_y(num_robot);

    /* Mobile robot poses and for next poses */
    Eigen::VectorXd cur_x(num_robot);
    Eigen::VectorXd cur_y(num_robot);
    Eigen::VectorXd obs_x(num_obs);
    Eigen::VectorXd obs_y(num_obs);

    /* Get swarm robot poses firstly */
    std::vector<std::vector<double> > current_robot_pose(swarm_robot_id.size());
    swarm_robot.getRobotPose(current_robot_pose);
    for(int i = 0; i < swarm_robot_id.size(); i++) {
        if(i < num_robot){
            cur_x(i) = current_robot_pose[i][0];
            cur_y(i) = current_robot_pose[i][1];
        }
        else{
            obs_x(i-num_robot) = current_robot_pose[i][0];
            obs_y(i-num_robot) = current_robot_pose[i][1];
        }
        
    }
    double obs_width = abs(obs_y(1) - obs_y(0));

    /* Generate Formation */
    swarm_robot.FormationGenerate(obs_width, interval, formation_x_list, formation_y_list);
    cout << formation_x_list[2] << endl;
    cout << formation_y_list[2] << endl;
    cout << "\033[33m>>>>>>>>>>>>>>>>>>>>>>>>>>>> To Formation Cross <<<<<<<<<<<<<<<<<<<<<<<<<\033[0m" << endl;

    /* Hungary */
    hungary_swarm.hungary(cur_x, cur_y, formation_x_list[0], formation_y_list[0]);
    for (int i = 0; i < num_robot; i++){
        cout << "Angent" << i + 1 << "-->"<< "Target" << hungary_swarm.assign[i] + 1 << endl;
        final_target_x(i) = formation_x_list[0](hungary_swarm.assign[i]);
        final_target_y(i) = formation_y_list[0](hungary_swarm.assign[i]);
    }
    /* FormationProcess */
    Eigen::Vector2d expected_point;
    expected_point(0) = cur_x.mean();
    expected_point(1) = cur_y.mean();
    swarm_robot.FormationProcess(expected_point, final_target_x, final_target_y, 
                                 k_v, MAX_V, MAX_W, conv_x, conv_y, conv_dis, swarm_vel);
    swarm_robot.stopRobot();
    cout << "\033[32m>>>>>>>>>>>>>>>>>>> Successfully! Press any key to continue<<<<<<<<<<<<<<<<<\033[0m" << endl;
    cin.get();


    cout << "\033[33m>>>>>>>>>>>>>>>>>>> Select a Formation <<<<<<<<<<<<<<<<<\033[0m" << endl;

    /* Get swarm robot poses */
    swarm_robot.getRobotPose(current_robot_pose);

    for(int i = 0; i < num_robot; i++) {
        cur_x(i) = current_robot_pose[i][0];
        cur_y(i) = current_robot_pose[i][1];
    }

    /* Formation Select */ 
    double final_assign[num_robot];
    double totalCost_min = 10000.0;
    int final_formation_index = 0;

    for (int i = 1; i < num_formation; i++) //TODO
    {
        hungary_swarm.hungary(cur_x, cur_y, formation_x_list[i], formation_y_list[i]);
        cout << "The total cost of formation \033[32m" << formation_name_list[i] 
             << "\033[0m is: \033[32m" << hungary_swarm.totalCost << "\033[0m" << endl;
        if (hungary_swarm.totalCost < totalCost_min)
        {
            totalCost_min = hungary_swarm.totalCost;
            final_formation_index = i;
            for (int j = 0; j < num_robot; j++){
                final_assign[j] = hungary_swarm.assign[j];
            }
        }
    }

    cout << "The final formation is: \033[32m" 
         << formation_name_list[final_formation_index] << "\033[0m" << endl;
    for (int i = 0; i < num_robot; i++)
    {
        cout << "Angent" << i + 1 << "-->"<< "Target" << final_assign[i] + 1 << endl;
        final_target_x(i) = formation_x_list[final_formation_index](final_assign[i]);
        final_target_y(i) = formation_y_list[final_formation_index](final_assign[i]);
    }

    cout << "press any key to continue" << endl;
    cin.get();

    expected_point(0) = cur_x.mean();
    expected_point(1) = cur_y.mean();
    swarm_robot.FormationProcess(expected_point, final_target_x, final_target_y, 
                                 k_v, MAX_V, MAX_W, conv_x, conv_y, conv_dis, swarm_vel);
    swarm_robot.stopRobot();                             
    cout << "\033[32m>>>>>>>>>>>>>>>>>>> Successfully! Press any key to continue<<<<<<<<<<<<<<<<<\033[0m" << endl;
    cin.get();


    cout << "\033[33m>>>>>>>>>>>>>>>>>>> Alias y to Obstacle in Selected Formation <<<<<<<<<<<<<<<<<\033[0m" << endl;
    expected_point(0) = obs_x.mean() - 1.2;
    expected_point(1) = obs_y.mean();
    swarm_robot.FormationProcess(expected_point, final_target_x, final_target_y, 
                                 k_v, MAX_V, MAX_W, conv_x, conv_y, conv_dis, swarm_vel);
    swarm_robot.stopRobot();                             
    cout << "\033[32m>>>>>>>>>>>>>>>>>>> Successfully! Press any key to continue<<<<<<<<<<<<<<<<<\033[0m" << endl;
    cin.get();


    cout << "\033[33m>>>>>>>>>>>>>>>>>>> Move to Right in Selected Formation <<<<<<<<<<<<<<<<<\033[0m" << endl;
    expected_point(0) = obs_x.mean() + 1.2;
    expected_point(1) = obs_y.mean();
    swarm_robot.FormationProcess(expected_point, final_target_x, final_target_y, 
                                 k_v, MAX_V, MAX_W, conv_x, conv_y, conv_dis, swarm_vel);
    swarm_robot.stopRobot();                             
    cout << "\033[32m>>>>>>>>>>>>>>>>>>> Successfully! Press any key to continue<<<<<<<<<<<<<<<<<\033[0m" << endl;
    cin.get();


    cout << "\033[33m>>>>>>>>>>>>>>>>>>> Switch to Init Formation <<<<<<<<<<<<<<<<<\033[0m" << endl;

    /* Get swarm robot poses */
    swarm_robot.getRobotPose(current_robot_pose);
    for(int i = 0; i < num_robot; i++) {
        cur_x(i) = current_robot_pose[i][0];
        cur_y(i) = current_robot_pose[i][1];
    }

    /* Hungary */
    hungary_swarm.hungary(cur_x, cur_y, formation_x_list[0], formation_y_list[0]);
    for (int i = 0; i < num_robot; i++){
        cout << "Angent" << i + 1 << "-->"<< "Target" << hungary_swarm.assign[i] + 1 << endl;
        final_target_x(i) = formation_x_list[0](hungary_swarm.assign[i]);
        final_target_y(i) = formation_y_list[0](hungary_swarm.assign[i]);
    }

    expected_point(0) = cur_x.mean();
    expected_point(1) = cur_y.mean();
    swarm_robot.FormationProcess(expected_point, final_target_x, final_target_y, 
                                 k_v, MAX_V, MAX_W, conv_x, conv_y, conv_dis, swarm_vel);

    /* Stop all robots */
    swarm_robot.stopRobot();

    ROS_INFO_STREAM("Succeed!");
    return 0;
}


/* Read parameter list */
void getParameterList(ros::NodeHandle nh){
    nh.getParam("conv_x", conv_x);
    nh.getParam("conv_y", conv_y);
    nh.getParam("intervalconv_dis", conv_dis);
    nh.getParam("intervalMAX_W", MAX_W);
    nh.getParam("intervalMIN_W", MIN_W);
    nh.getParam("intervalMAX_V", MAX_V);
    nh.getParam("intervalMIN_V", MIN_V);
    nh.getParam("intervalk_w", k_w);
    nh.getParam("intervalk_v", k_v);
    nh.getParam("intervalweight_repel", weight_repel);
    nh.getParam("intervalweight_attract", weight_attract);
    nh.getParam("intervaldis_repel", dis_repel);
    nh.getParam("intervaldis_attract", dis_attract);
    nh.getParam("intervalswarm_vel", swarm_vel);
    nh.getParam("intervalinterval", interval);

    ROS_INFO("Parameters list reading finished!");
}